'''
author:        wangchenyang <cy-wang21@mails.tsinghua.edu.cn>
date:          2025-03-21
Copyright © Department of Physics, Tsinghua University. All rights reserved

Investigate the HN model without the GDSE.
'''

import numpy as np
from numpy import pi
import pyvista as pv
import poly_tools as pt
import matplotlib.pyplot as plt
from palette_manager import get_palette

COLOR = get_palette("matplotlib")

def get_ChP(gammax, gammay, Jx, Jy, alpha):
    coeffs = pt.CScalarVec(
        [1, -gammax * Jx, 
         -Jx.conjugate() / gammax, 
         -gammay * Jy, 
         -Jy.conjugate() / gammay]
    )
    degs = pt.CLaurentIndexVec([
        # E, beta1, beta2
        1, 0, 0,
        0, -1, alpha,
        0, 1, -alpha,
        0, 0, -1,
        0, 0, 1
    ])
    chp = pt.CLaurent(3)
    chp.set_Laurent_by_terms(coeffs, degs)
    return chp


def get_ChP_eq(chp: pt.CLaurent, E, beta1):
    chp_1d = chp.partial_eval(
        pt.CScalarVec([E, beta1]),
        pt.CIndexVec([0, 1]),
        pt.CIndexVec([2])
    )
    M = chp_1d.denom_orders[0]
    coeffs = pt.CScalarVec([])
    degs = pt.CIndexVec([])
    chp_1d.num.batch_get_data(coeffs, degs)

    # Get numpy array
    max_deg = max(degs)
    eq_np = np.zeros(max_deg + 1, dtype=complex)
    for i in range(len(coeffs)):
        eq_np[max_deg - degs[i]] = coeffs[i]
    return eq_np, M


def plot_mGBZ():
    Jx = 1
    Jy = 0.5 + 0.2j
    alpha = 2
    gammax = 0.8
    gammay = 0.7

    r = gammax * (gammay ** alpha) * 1.0

    chp = get_ChP(gammax, gammay, Jx, Jy, alpha)
    E = 0.2
    theta1 = np.linspace(-pi, pi, 1000)
    beta1 = r * np.exp(1j * theta1)
    beta2 = np.zeros((2, len(beta1)), dtype=complex)
    picked_roots = []
    for idx, curr_beta1 in enumerate(beta1):
        eq_np, M = get_ChP_eq(chp, E, curr_beta1)
        all_roots = np.roots(eq_np)
        all_roots = sorted(all_roots, key=lambda x: abs(x))
        beta2[0, idx] = all_roots[M - 1]
        beta2[1, idx] = all_roots[M]
        picked_roots += [(curr_beta1, beta2) for beta2 in all_roots if abs(abs(beta2) - gammay) / gammay < 0.4]
    # Plot
    plt.figure()
    plt.plot(theta1, np.abs(beta2[0,:]))
    plt.plot(theta1, np.abs(beta2[1,:]))
    plt.plot([-pi, pi], [gammay, gammay], '--')
    plt.xlim([-pi, pi])

    picked_roots = np.asarray(picked_roots)
    plt.figure()
    plt.plot(np.log(picked_roots[:, 1]).imag, np.log(picked_roots[:, 0]).imag, '.')
    plt.xlim([-pi, pi])
    plt.ylim([-pi, pi])
    plt.show()

    pl = pv.Plotter()
    pl.add_points(
        np.column_stack([
            picked_roots[:, 1].real,
            picked_roots[:, 1].imag,
            np.log(picked_roots[:, 0]).imag
        ]),
        color=COLOR['ol']
    )
    cld = pv.Cylinder(radius=gammay, height=2*pi, direction=(0, 0, 1))
    pl.add_mesh(cld, color=COLOR['bf'], opacity=0.5)
    pl.show()

if __name__ == '__main__':
    plot_mGBZ() 
